import grpc

import server.service.pb.model_server_pb2 as pb
from server.service.pb.model_server_pb2_grpc import ModelServerStub
from server.service.utils.pb_helper import *


class ModelServerClient:
    def __init__(self, ip, port=50051):
        # TODO(更改了消息大小的限制，否则一些很大的数据集（例如few-nerd数据集）的预测结果无法返回)
        self.channel = grpc.insecure_channel(f'{ip}:{port}', options=[
            ('grpc.max_receive_message_length', 10 * 1024 * 1024),  # 设置为10MB
            ('grpc.max_send_message_length', 10 * 1024 * 1024),  # 设置为10MB
        ])
        self.stub = ModelServerStub(self.channel)

    def create_experiment(self, task, model, dataset=None, experiment_id=None, project_id=None, version=None,
                          model_config=None,
                          optimizer_config=None, data=None, label_list=None):
        experiment = PbHelper().create_experiment(task=task, model=model, dataset=dataset, experiment_id=experiment_id,
                                                  project_id=project_id, version=version,
                                                  model_config=model_config,
                                                  optimizer_config=optimizer_config, data=data, label_list=label_list)
        response = self.stub.create_experiment(experiment)
        return response

    def predict(self, text, task, model, dataset=None, version=None, experiment_id=None, project_id=None,
                model_config=None, label_list=None):
        assert isinstance(text, list)
        request = PbHelper().create_predict_request(text=text, task=task, model=model, dataset=dataset,
                                    version=version,
                                    experiment_id=experiment_id,project_id=project_id,
                                    model_config=model_config,label_list=label_list)
        response = self.stub.predict(request)

        return response

    def texttool(self, model, label_list, train_data, test_data, task, version, project_id,
                 model_config=None,
                 optimizer_config=None, acquire=10):
        texttool_request = create_texttool_request(model=model, label_list=label_list,
                                                   train_data=train_data, test_data=test_data,
                                                   task=task,
                                                   version=version, project_id=project_id,
                                                   model_config=model_config,
                                                   optimizer_config=optimizer_config,
                                                   acquire=acquire)
        response = self.stub.texttool(texttool_request)

        return response

    def close(self):
        self.channel.close()
